#include "thread_pool.h"
#include <stdlib.h>

static void task_offer(thread_pool_t* pool, thread_pool_task_t* new_task)
{
    thread_pool_task_queue_t* queue = pool->queue;
    if (queue == NULL)
    {
        //TODO unexpected
        return;
    }
    pthread_mutex_lock(&queue->lock);
    thread_pool_task_t* task = queue->tail;
    if (task == NULL) {
        queue->head = new_task;
        queue->tail = new_task;
    } else {
        task->next = new_task;
        queue->tail = new_task;
    }
    queue->count++;
    pthread_mutex_unlock(&queue->lock);
}

static thread_pool_task_t* task_poll(thread_pool_t* pool)
{
    thread_pool_task_queue_t* queue = pool->queue;
    if (queue == NULL)
    {
        //TODO unexpected
        return NULL;
    }
    pthread_mutex_lock(&queue->lock);
    thread_pool_task_t* task = queue->head;
    if (task != NULL)
    {
        queue->head = task->next;
        if (task->next == NULL)
        {
            queue->tail = NULL;
        }
        queue->count--;
    }
    pthread_mutex_unlock(&queue->lock);
    return task;
}

static thread_pool_task_queue_t* task_queue_create()
{
    thread_pool_task_queue_t* queue = (thread_pool_task_queue_t*)malloc(sizeof(thread_pool_task_queue_t));
    if (queue == NULL)
    {
        //Failed to allocate memory for task queue
        return NULL;
    }
    queue->head = NULL;
    queue->tail = NULL;
    queue->count = 0;
    pthread_mutex_init(&queue->lock, NULL);

    return queue;
}

static int task_queue_count(thread_pool_t* pool)
{
    if (pool->queue == NULL)
    {
        return 0;
    }
    return pool->queue->count;
}

static void task_queue_clean(thread_pool_task_queue_t* queue)
{
    thread_pool_task_t* task = queue->head;
    thread_pool_task_t* next;
    while (task != NULL)
    {
        next = task->next;
        free(task);
        task = next;
    }
    queue->head = NULL;
    queue->tail = NULL;
    queue->count = 0;
}

static void task_queue_destroy(thread_pool_t* pool)
{
    if (pool->queue != NULL)
    {
        task_queue_clean(pool->queue);
        pthread_mutex_destroy(&pool->queue->lock);
        free(pool->queue);
        pool->queue = NULL;
    }
}

static void* thread_entry(void* arg)
{
    thread_pool_t* pool = (thread_pool_t*)arg;
    thread_pool_task_t* task;

    while (1)
    {
        pthread_mutex_lock(&pool->lock);
        while (task_queue_count(pool) == 0 && pool->keepalive)
        {
            pthread_cond_wait(&pool->cond, &pool->lock);
        }
        pthread_mutex_unlock(&pool->lock);

        task = task_poll(pool);
        if (task != NULL)
        {
            task->func(task->arg);
            free(task);
            continue;
        }

        if (task == NULL && !pool->keepalive)
            break;
    }

    return NULL;
}

thread_pool_t* thread_pool_create(int thread_count)
{
    thread_pool_t* pool = (thread_pool_t*)malloc(sizeof(thread_pool_t));
    if (pool == NULL)
    {
        // Failed to allocate memory for thread pool
        return NULL;
    }

    pool->thread_count = thread_count;
    pool->keepalive = 1;

    pool->queue = task_queue_create();

    if (pthread_mutex_init(&pool->lock, NULL) != 0
        || pthread_cond_init(&pool->cond, NULL) != 0)
    {
        //Failed to initialize mutex or condition variable
        goto err_cond;
    }

    pool->threads = (pthread_t*)malloc(thread_count * sizeof(pthread_t));
    if (pool->threads == NULL)
    {
        //Failed to allocate memory for threads
        goto err_threads;
    }

    for (int i = 0; i < thread_count; i++)
    {
        if (pthread_create(&pool->threads[i], NULL, thread_entry, pool) != 0)
        {
            //Failed to create thread
            for (int j = 0; j < i; j--) {
                pthread_cancel(pool->threads[j]);
            }
            goto err_thread;
        }
    }

    return pool;

err_thread:
    free(pool->threads);

err_threads:
    pthread_mutex_destroy(&pool->lock);
    pthread_cond_destroy(&pool->cond);

err_cond:
    free(pool);
    return NULL;
}

void thread_pool_destroy(thread_pool_t* pool)
{
    if (pool == NULL)
    {
        return;
    }

    pthread_mutex_lock(&pool->lock);
    pool->keepalive = 0;

    pthread_cond_broadcast(&pool->cond);
    pthread_mutex_unlock(&pool->lock);

    for (int i = 0; i < pool->thread_count; i++) {
        pthread_join(pool->threads[i], NULL);
    }

    free(pool->threads);
    task_queue_destroy(pool);
    pthread_mutex_destroy(&pool->lock);
    pthread_cond_destroy(&pool->cond);
    free(pool);
}

void thread_pool_submit(thread_pool_t* pool, void(*func)(void*), void* arg)
{
    if (pool == NULL && func == NULL)
    {
        return;
    }

    if (!pool->keepalive)
    {
        return;
    }

    thread_pool_task_t* new_task = (thread_pool_task_t*)malloc(sizeof(thread_pool_task_t));
    if (new_task == NULL)
    {
        //Failed to allocate memory for task
        return;
    }

    new_task->func = func;
    new_task->arg = arg;
    new_task->next = NULL;

    task_offer(pool, new_task);

    pthread_mutex_lock(&pool->lock);
    pthread_cond_signal(&pool->cond);
    pthread_mutex_unlock(&pool->lock);
}

